package main

import (
	"github.com/gin-gonic/gin"
	ratelimit2 "github.com/juju/ratelimit"
	ratelimit1 "go.uber.org/ratelimit"
	"net/http"
	"time"
)

func PingHandler(c *gin.Context) {
	c.String(200, "通了！")

}

func HeiHandler(c *gin.Context) {
	c.String(200, "嘿哈")
}

func main() {

	r := gin.Default()
	r.Use(RateLimit2())
	r.GET("/ping", PingHandler)
	r.GET("/hei", HeiHandler)
	r.Run(":9999")

}

// 基于漏桶的限流策略, 中间件
func RateLimit1() func(c *gin.Context) {
	// 返回一个 func(c *gin.Context类型的函数
	r1 := ratelimit1.New(100)
	return func(c *gin.Context) {
		// 生成一个限流器
		// Take() 返回的是需要等待的时间， 如果需要等待的时间 - 当前时间 > 0 表示还没过需要等待的时间
		if r1.Take().Sub(time.Now()) > 0 {
			c.String(200, "请求太过频繁！")
			c.Abort()
			return


		}
		c.Next()

	}
}
// 基于令牌桶的限流方式
/*
适合处理突发的大面积请求
不需要特意实现令牌桶存令牌取出令牌的操作
 */
func RateLimit2() func(c *gin.Context) {
	// 返回一个 func(c *gin.Context类型的函数
	bucket := ratelimit2.NewBucket(time.Second * 2, 1)
	return func(c *gin.Context) {
		// 如果取不到令牌就中断本次请求返回 rate limit...
		// take可以欠账 takeavilable 不可以欠账
		if bucket.TakeAvailable(1) < 1 { // 从桶中拿取一个令牌，如果桶中的令牌数小于1，则触发限流
			c.String(http.StatusOK, "rate limit...")
			c.Abort()
			return
		}
		c.Next()

	}
}
